//This function performs clustering of points in R^3.
//usage example : a = rand(3,10,10); Kmeans3(a,4);

#include "mex.h"
#include <string.h>
#include <cmath>
#include <cstdlib>
#include <ctime>
#include<iostream>
using namespace std;
#define INF 1e9
#define MAXSTEPS 100
inline double dist(const double*a, double* c, int i, int j)
{
	double x = a[3 * i] - c[3 * j] ;
	double y = a[3 * i + 1] - c[3 * j + 1];
	double z = a[3 * i + 2] - c[3 * j + 2];
	x = x > 0 ? x : -x;
	y = y > 0 ? y : -y;
	z = z > 0 ? z : -z;
	return x + y + z;
}

int findCenter3(const double* a, double* c, int i, int k)
{
	double min = INF;
	int t = -1;
	for (int j = 0; j < k; j++)
	{
		if ( dist(a, c, i, j) < min )
		{
			min = dist(a, c, i, j);
			t = j;
		}
	}
	return t;
}

void kmeans3(const double* a, int n, int k, int*b, double* c)
{
	int* p = new int [n];
	memset(p, 0, sizeof(int)*n);	
	for (int i = 0; i < k; i++)
	{
		int x = int(n*(double)rand()/RAND_MAX);
		memcpy(c+i*3,a+x*3,sizeof(double)*3);
	}
	bool flag = false;
	int step = 0; 
	bool flagCenter = true;
	while (!flag)
	{

		flagCenter = true;
		int err = 0;
		int* count = new int [k];
		double* s = new double [3 * k];
		memset(count, 0, sizeof(int)*k);
		memset(s, 0, sizeof(double)*3*k);
		for (int i = 0; i < n; i++)
		{
			int t = findCenter3(a, c, i, k); //?
			s[3*t] += a[3*i];
			s[3*t+1] += a[3*i+1];
			s[3*t+2] += a[3*i+2];
			count[t]++;
			err += p[i] - t > 0 ? p[i] - t : t - p[i];
			p[i] = t;
		}
		for (int i = 0; i < k; i++)
		{
			if (count[i] == 0) 
			{
				for (int i = 0; i < k; i++)
				{
					int x = int(n*(double)rand()/RAND_MAX);
					memcpy(c+i*3,a+x*3,sizeof(double)*3);
				}
				break;
			}
			else
			{
				for (int i = 0; i < k; i++)
				{
					c[3*i] = s[3*i] / count[i];
					c[3*i+1] = s[3*i+1] / count[i];
					c[3*i+2] = s[3*i+2] / count[i];
				}
				if ( err == 0 || step == MAXSTEPS) flag = true;
			}
		}
		step++;	
	}
	memcpy(b, p, sizeof(int)*n);
	delete [] p;
}

double opt(const double* a, int n, int k, int* p, double* c)
{
	double s = 0;
	double x1,x2,x3;
	for (int i = 0; i < n; i++)
	{
		x1 = a[i*3]-c[p[i]*3];
		x2 = a[i*3 + 1] - c[p[i]*3+1];
		x3 = a[i*3 + 2] - c[p[i]*3+2];
		s += x1*x1+x2*x2+x3*x3;
	}
	return s;
}

//int main()
//{
//	srand((time(0)));
//	int n = 200000;
//	int k = 3;
//	double* a= new double [n*3];
//
//	int *p = new int [n];
//	memset(p, 0, sizeof(int)*n);
//	for (int i = 0;i < n*3; i++)
//		a[i] = (double)rand()/RAND_MAX;
//	double min = INF;
//	int* tmpb = new int [n];
//	double* tmpc = new double [k*3];
//	int* b = new int [n];
//	double* c = new double [k*3];
//	for (int i = 0 ; i < ITER; i++)
//	{	
//		kmeans3(a, n, k, tmpb, tmpc);
//		double optVal = opt(a, n, k, tmpb, tmpc);
//		cout<<min<<endl;
//		if ( optVal < min)
//		{
//			min = optVal;
//			memcpy(b, tmpb, sizeof(int)*n);
//			memcpy(c, tmpc, sizeof(double)*k*3);
//		}
//	}
//	system("pause");
//	return 0;
//}

void mexFunction(int nlhs, mxArray* plhs[], int nrhs, mxArray* prhs[])
{
	if ( nrhs < 2)
		mexErrMsgTxt("2 input arguments is required.");
	else if (nrhs > 3)
		mexErrMsgTxt("Too many input arguments.");
	else if( nlhs > 2)
		mexErrMsgTxt("Too many output arguments.");

	int ITER = 10;
	if ( nrhs == 3)
	{
		ITER = (int)mxGetScalar(prhs[2]);
	}
	if( ITER < 1)
		mexErrMsgTxt("Min iteration number must be larger than 1.");
	int m = mxGetM(prhs[0]);
	if ( m != 3 ) 
		mexErrMsgTxt("First dimension must be 3.");
	int n = mxGetN(prhs[0]);
	int k = (int)mxGetScalar(prhs[1]);
	double* a = mxGetPr(prhs[0]);
	plhs[0] = mxCreateNumericMatrix(3,k, mxDOUBLE_CLASS, mxREAL);
	const int* D = mxGetDimensions(prhs[0]);
	const int N = mxGetNumberOfDimensions(prhs[0]);
	int* DIM = new int [N-1];
	for (int i = 0; i < N-1; i++)
	{
		DIM[i] = D[i+1];
	}
	plhs[1] = mxCreateNumericArray(N - 1, DIM, mxINT32_CLASS, mxREAL);
	int* b = (int*)mxGetPr(plhs[1]);
	double* c = mxGetPr(plhs[0]);

	double min = INF;
	int* tmpb = new int [n];
	double* tmpc = new double [k*3];
	for (int i = 0 ; i < ITER; i++)
	{	
		kmeans3(a, n, k, tmpb, tmpc);
		double optVal = opt(a, n, k, tmpb, tmpc);
		if ( optVal < min)
		{
			min = optVal;
			memcpy(b, tmpb, sizeof(int)*n);
			memcpy(c, tmpc, sizeof(double)*k*3);
		}
	}
}

